package utils

import (
	"sync"
	"time"
)

const (
	SEQUENCE_BIT       = 12
	MACHINE_BIT        = 5
	DATACENTER_BIT     = 5
	MAX_SEQUENCE       = 4095
	MACHINE_LEFT       = 12
	DATACENTER_LEFT    = 17
	TIMESTAMP_LEFT     = 22
	MAX_MACHINE_NUM    = 31
	MAX_DATACENTER_NUM = 31
)

type SnowflakeKeyGenerator struct {
	mutex        sync.Mutex
	sequence     int64
	lastStamp    int64
	datacenterId int64
	machineId    int64
	startStamp   int64
}

func NewSnowflakeKeyGenerator(datacenterId, machineId int64) *SnowflakeKeyGenerator {
	if datacenterId > MAX_DATACENTER_NUM || datacenterId < 0 {
		panic("datacenterId can't be greater than 31 or less than 0")
	}
	if machineId > MAX_MACHINE_NUM || machineId < 0 {
		panic("machineId can't be greater than 31 or less than 0")
	}

	return &SnowflakeKeyGenerator{
		datacenterId: datacenterId,
		machineId:    machineId,
		startStamp:   time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC).UnixMilli(),
		lastStamp:    -1,
	}
}

func (s *SnowflakeKeyGenerator) GenerateKey() int64 {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	currStamp := s.getNewTimeStamp()
	if currStamp < s.lastStamp {
		panic("Clock moved backwards. Refusing to generate id")
	}

	if currStamp == s.lastStamp {
		s.sequence = (s.sequence + 1) & MAX_SEQUENCE
		if s.sequence == 0 {
			currStamp = s.getNextMill()
		}
	} else {
		s.sequence = 0
	}

	s.lastStamp = currStamp
	return (currStamp-s.startStamp)<<TIMESTAMP_LEFT |
		(s.datacenterId << DATACENTER_LEFT) |
		(s.machineId << MACHINE_LEFT) |
		s.sequence
}

func (s *SnowflakeKeyGenerator) getNewTimeStamp() int64 {
	return time.Now().UnixMilli()
}

func (s *SnowflakeKeyGenerator) getNextMill() int64 {
	curr := s.getNewTimeStamp()
	for curr <= s.lastStamp {
		curr = s.getNewTimeStamp()
	}
	return curr
}
